﻿#region Copyright information
// 
// Copyright © 2005-2013 Yongkee Cho. All rights reserved.
// 
// This code is a part of the Biological Object Library and governed under the terms of the
// GNU Lesser General  Public License (LGPL) version 2.1 which accompanies this distribution.
// For more information on the LGPL, please visit http://bol.codeplex.com/license.
// 
// - Filename: DirichletDistribution.cs
// - Author: Yongkee Cho
// - Email: yongkeecho@gmail.com
// - Date Created: 2012-09-06 11:39 AM
// - Last Modified: 2013-01-25 3:59 PM
// 
#endregion
using System;
using System.Collections.Generic;
using System.Linq;
using BOL.Maths.Functions;
using BOL.Algorithms.Optimization;

namespace BOL.Maths.Distributions
{
    /// <summary>
    /// Represents a Dirichlet distribution with parameter vector alpha.
    /// <see>http://en.wikipedia.org/wiki/Dirichlet_distribution</see>
    /// </summary>
    public sealed class DirichletDistribution : IMultivariateDistribution<double>, IEquatable<DirichletDistribution>
    {
        #region Private variables

        private double[] _alpha; // hyperparameter alpha
        private int _k; // length of alpha
        private double _alpha0; // sum of alpha, commonly referred to precision or concentration parameter

        #endregion

        #region Public properties

        /// <summary>Gets or sets the alpha of a Dirichlet distribution.</summary>
        public double[] Alpha
        {
            get { return _alpha; }
            set
            {
                _alpha = value;
                _k = value.Length;
                _alpha0 = value.Sum();
            }
        }

        /// <summary>Gets the sum of alpha of a Dirichlet distribution.</summary>
        public double Alpha0 { get { return _alpha0; } }

        /// <summary>Gets the length of alpha of a Dirichlet distribution.</summary>
        public int K { get { return _k; } }

        /// <summary>Gets the range of a Dirichlet distribution.</summary>
        public IRange<double> Domain { get { return new Range<double>(0, 1); } }

        /// <summary>Gets the mean of a Dirichlet distribution.</summary>
        public double[] Mean { get { return _alpha.Select(alpha => alpha / _alpha0).ToArray(); } }

        /// <summary>Gets the mode of a Dirichlet distribution.</summary>
        public double[] Mode { get { return _alpha.Select(alpha => (alpha - 1) / (_alpha0 - _k)).ToArray(); } }

        /// <summary>Gets the covariance of a Dirichlet distribution.</summary>
        public double[][] Covariance
        {
            get
            {
                var cov = new double[_k][];
                var denominator = _alpha0 * _alpha0 * (_alpha0 + 1);

                for (var i = 0; i < _k; i++)
                {
                    cov[i] = new double[_k];
                    for (var j = 0; j < _k; j++)
                        cov[i][j] = i == j ? 
                            _alpha[i] * (_alpha0 - _alpha[i]) / denominator : 
                            -_alpha[i] * _alpha[j] / denominator;
                }
                return cov;
            }
        }

        /// <summary>Gets the entropy of a Dirichlet distribution.</summary>
        /// <see cref="http://en.wikipedia.org/wiki/Dirichlet_distribution#Entropy"/>
        public double Entropy
        {
            get
            {
                var term1 = 1.0;
                var term2 = (_alpha0 - _k) * DistributionFunctions.DiGamma(_alpha0);
                var term3 = 0.0;

                for (var i = 0; i < _k; i++)
                {
                    term1 *= DistributionFunctions.Gamma(_alpha[i]);
                    term3 += (_alpha[i] - 1) * DistributionFunctions.DiGamma(_alpha[i]);
                }
                term1 /= DistributionFunctions.Gamma(_alpha0);
                term1 = Math.Log(term1);

                return term1 + term2 - term3;
            }
        }

        #endregion

        #region Constructor

        /// <summary>
        /// Instantiates a Dirichlet distribution.
        /// </summary>
        /// <param name="alpha">parameter vector</param>
        public DirichletDistribution(double[] alpha)
        {
            if (alpha == null)
                throw new ArgumentNullException("alpha");
            if (alpha.Any(a => a < 0))
                throw new ArgumentOutOfRangeException("alpha", "alpha are expected to be positive.");
            
            _alpha = alpha;
        }

        #endregion

        #region ICloneable implementation

        public DirichletDistribution Clone()
        {
            return new DirichletDistribution(_alpha.Select(a => a).ToArray());
        }

        object ICloneable.Clone()
        {
            return Clone();
        }
        
        #endregion

        #region Public methods

        /// <summary>
        /// Returns the Dirichlet density given value.
        /// </summary>
        /// <param name="value">proportion vector</param>
        /// <returns>Dirichlet density</returns>
        public double Pdf(double[] value)
        {
            if (value == null)
                throw new ArgumentNullException("value");

            if (value.Length != _k)
                throw new ArgumentException("The length of value vector is expected to be k.", "value");

            if (value.Any(v => v < 0))
                throw new ArgumentOutOfRangeException("value", "value is expected to be greater than zero.");

            if (Math.Sqrt(value.Sum() - 1.0) > 1e-10)
                throw new ArgumentException("x_1 + x_2 + ... + x_k is expected to be one.", "value");

            // normalizing constant is beta/gamma function of alpha
            var beta = _alpha.Aggregate(1.0, (b, alphaI) => b * DistributionFunctions.Gamma(alphaI)) / DistributionFunctions.Gamma(_alpha0);
            return _alpha.Zip(value, (alphaI, valueI) => Math.Pow(valueI, alphaI - 1))
                .Aggregate(1.0, (pdf, zip) => pdf * zip)
                / beta;
        }

        /// <summary>
        /// Returns the cummulative Dirichlet density given value.
        /// </summary>
        /// <param name="value">proportion vector</param>
        /// <returns>cummulative Dirichlet density</returns>
        public double Cdf(double[] value)
        {
            throw new Exception("Cdf is undefined for Dirichlet distribution.");
        }

        /// <summary>
        /// Returns the value of the Dirichlet distribution for given probability.
        /// </summary>
        /// <param name="p">probability</param>
        /// <returns></returns>
        public double[] Quantile(double p)
        {
            throw new NotImplementedException();
        }

        /// <summary>
        /// Returns a random value from the Dirichlet distribution.
        /// <see>http://en.wikipedia.org/wiki/Dirichlet_distribution#Gamma_distribution</see>
        /// </summary>
        /// <returns>random value</returns>
        public double[] Sample(Random r)
        {
            var sample = _alpha.Select(alpha => new GammaDistribution(alpha, 1).Sample(r)).ToList();
            var sum = sample.Sum();
            return sample.Select(s => s / sum).ToArray();
        }

        /// <summary>
        /// Estimates parameters of the Dirichlet distribution using maximum likelihood.
        /// Given a set of observed multinomial data, D = {p1, p2, ..., pn},
        /// the parameters for a Dirichlet distribution can be estimated 
        /// by maximizing the log-likelihood function of the data
        /// <see>http://www.cs.cmu.edu/~jch1/research/dirichlet/dirichlet.pdf</see>
        /// </summary>
        public void MaximumLikelihoodEstimate(IEnumerable<double[]> source)
        {
            //if (source == null)
            //    throw new ArgumentNullException("source");

            //var list = source.ToList();
            //var k = list[0].Length;
            //var alpha = new double[k];
            //foreach (var l in list)
            //    for (var i = 0; i < k; i++)
            //        alpha[i] += l[i];
            //var dll = new DirichletLogLikelihood(list, alpha);
            //var nr = new NewtonRaphson(dll, 100, 1e-5, true);
            //nr.Optimize();
            //return new DirichletDistribution(dll.Parameters);

            throw new NotImplementedException();
        }

        /// <summary>
        /// Estimates parameters of the Dirichlet distribution using Bayesian methods.
        /// </summary>
        public void BayesianEstimate(IEnumerable<double[]> source, params dynamic[] priorParameters)
        {
            throw new NotImplementedException();
        }

        #region Dirichlet log-likelihood function

        /// <summary>
        /// Implements a Dirichlet log-likelihood function.
        /// </summary>
        private class DirichletLogLikelihood : ObjectiveFunctionBase
        {
            #region Private variables

            private readonly int _length; // number of instances
            private readonly double[] _logSum; // log pK = sigma_i(log(p_ik)) / N

            #endregion 

            #region Constructor

            public DirichletLogLikelihood(IEnumerable<double[]> data, double[] alpha)
                : base(alpha)
            {
                if (data == null)
                    throw new ArgumentNullException("data");

                //current = alpha.Select(a => a).ToArray();
                var k = alpha.Length;
                var list = data.ToList();
                _length = list.Count;

                _logSum = new double[k];
                for (var j = 0; j < k; j++)
                {
                    var logSum = 0.0;
                    for (var i = 0; i < _length; i++)
                        if (list[i][j] > 0) // prevent from -infinity
                            logSum += Math.Log(list[i][j]);
                    _logSum[j] = logSum / _length;
                }
            }

            #endregion

            #region Public methods

            public override double Value(double[] alpha)
            {
                if (alpha == null)
                    throw new ArgumentNullException("alpha");

                if (alpha.Length != NumberOfParameters)
                    throw new ArgumentException("The length of alpha mismatches.");

                var alpha0 = alpha.Sum();
                var value = DistributionFunctions.LogGamma(alpha0);

                for (var i = 0; i < NumberOfParameters; i++)
                    value -= DistributionFunctions.LogGamma(alpha[i]) + (alpha[i] - 1) * _logSum[i];
                value *= _length;

                return value;
            }

            public override double[] Gradient(double[] alpha)
            {
                if (alpha == null)
                    throw new ArgumentNullException("alpha");

                if (alpha.Length != NumberOfParameters)
                    throw new ArgumentException("The length of alpha mismatches.");

                var gradient = new double[NumberOfParameters];
                var alpha0 = alpha.Sum();
                var psi0 = DistributionFunctions.DiGamma(alpha0);

                for (var i = 0; i < NumberOfParameters; i++)
                {
                    var psi = DistributionFunctions.DiGamma(alpha[i]);
                    gradient[i] = _length * (psi0 - psi + _logSum[i]);
                }

                return gradient;
            }

            public override double[][] Hessian(double[] alpha)
            {
                if (alpha == null)
                    throw new ArgumentNullException("alpha");

                if (alpha.Length != NumberOfParameters)
                    throw new ArgumentException("The length of alpha mismatches.");

                var hessian = new double[NumberOfParameters][];
                var alpha0 = alpha.Sum();
                var psi0 = DistributionFunctions.TriGamma(alpha0);

                for (var i = 0; i < NumberOfParameters; i++)
                {
                    hessian[i] = new double[NumberOfParameters];

                    for (var j = 0; j < NumberOfParameters; j++)
                    {
                        if (i != j)
                            hessian[i][j] = _length * psi0;
                        else
                        {
                            var psi = DistributionFunctions.TriGamma(alpha[i]);
                            hessian[i][j] = _length * (psi0 - psi);
                        }
                    }
                }

                return hessian;
            }

            #endregion
        }

        #endregion

        #endregion

        #region IEquatable<T> implementation

        public bool Equals(DirichletDistribution other)
        {
            return _alpha.SequenceEqual(other._alpha);
        }

        #endregion

        #region Object overriden

        public override int GetHashCode()
        {
            return Alpha.Aggregate(1, (hashCode, alpha) => hashCode ^ alpha.GetHashCode());
        }

        public override bool Equals(object other)
        {
            if (other == null)
                throw new ArgumentNullException("other");

            if (!(other is DirichletDistribution))
                throw new InvalidCastException("The 'other' argument is not a DirichletDistribution object.");
            
            return Equals(other as DirichletDistribution);
        }

        public override string ToString()
        {
            return String.Format("Dirichlet(alpha = {{{0}}})", String.Join(",", _alpha));
        }

        #endregion
    }
}
